from langchain_openai import ChatOpenAI
from PIL import Image
from json_repair import repair_json

from ..config import settings
from ..utils.utils import image_to_b64str



class QwenVL(object):
    def __init__(self):
        self.vl = ChatOpenAI(
            base_url=settings.VL_URI,
            model=settings.VL_MODEL
        )
        self.system_prompt = """
你是Qwen-VL模型助手，接收图片后，请先对图像进行分类识别（文本类/表格类/流程图/公式/其他图像），然后再按以下要求进行处理。

## **处理规则**
- **文本类图像**：需精准识别所有可见文本，完整保留原始排版结构（如段落缩进、对齐方式、字体层级等格式），仅输出纯转录内容，不添加任何主观评论或解析。
- **表格类图像**：需识别表格行列结构（含合并单元格、跨列等复杂格式），提取单元格内文字信息，以标准HTML表格语法输出，你需要提取所有的详细信息，并确保行列对应关系与原图一致。
- **流程图**：需解析图形元素（矩形/菱形/箭头等）逻辑关系，转化为Mermaid语法的Markdown文本，包含节点定义、连接关系及箭头流向说明，你需要提取所有的详细信息，并保留原图流程逻辑。
- **公式**：需识别图片中的公式为有效的LaTex格式文本，并使得Latex文本的渲染效果与图片所示一致。
- **其他图像**：需从客观视角描述图像内容，包含主体元素、色彩基调、场景特征、关键细节（如尺寸比例、位置关系等），若存在可识别文本需同步标注。

## **结构化输出模板**
请严格按照以下 JSON 格式输出结果，禁止添加任何自然语言描述、注释或多余内容，确保格式完全符合 JSON 规范（如使用双引号、正确的逗号分隔、字段名准确等）。
```json
{
  "class": "<图像的分类，严格限定为以下枚举值之一：文本/表格/流程/公式/其他>", 
  "text": "<按处理规则进行处理后的文本内容。若为文本类图像，则直接输出转录内容；若为表格类图像，则输出html语法的表格内容；若为流程图，则输出Mermaid语法的Markdown格式文本；若为公式类图像，则输出Latex语法文本；若为其他图像，则输出图像的文本描述内容。>"
}
```
"""

    def invoke(self, image: Image.Image) -> dict:
        image_base64 = image_to_b64str(image)
        messages = [
            {
                'role': 'user',
                'content': [
                    {
                        "type": "text",
                        "text": self.system_prompt
                    },
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": f"data:image/jpeg;base64,{image_base64}"
                        }
                    }
                ]
            }
        ]
        resp = self.vl.invoke(messages)
        data = repair_json(resp.content, return_objects=True)
        if not isinstance(data, dict) or 'text' not in data:
            data = {'class': '其他', 'text': resp.content}
        return data

if __name__ == '__main__':
    qwen = QwenVL()
    image = Image.open(r'D:\workprojects\marker_patch\data\images\formula-2.png')
    result = qwen.invoke(image)
    print(result['text'])